import SimpleITK as sitk
import glob
import numpy as np
from PIL import Image
import nibabel as nib
from pathlib import Path
import os

import cv2

data_route = r'../segmentation_data'
BASEDIR = os.path.normpath(r'F:\黄陶然')
files = Path(BASEDIR).rglob('*2_image.nii.gz')

for filename in files:
    print(filename)
    file = str(filename).split('\\')[4][:-7]
    route = os.path.join(data_route, file)
    if not os.path.isdir(route):
        os.mkdir(route)
    imageNii = nib.load(str(filename))
    image = imageNii.get_fdata()
    if len(image.shape) > 3:
        image_shape = image.shape[-3]
    else:
        image_shape = image.shape[-1]
    for i in range(image_shape):
        # tempResult = image[:, :, i]
        print(image.shape)
        if len(image.shape) > 3:
            silce = np.array([image[:, :, i, 0, 0]]).transpose((-1, 1, 0))
        else:
            silce = np.array([image[:, :, i]]).transpose((-1, 1, 0))
        new_image = nib.Nifti1Image(silce, np.eye(4))
        nib.save(new_image, f'{route}/{file}_{"%04d" % i}_0000.nii.gz')
        print(f'{route}/{file}_{"%04d" % i}.nii.gz')
